
#
# A Production Rule System implemented in Python.
#
# (C) Copyright 2006-2008, Johan Lindberg
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 59 Temple
# Place, Suite 330, Boston, MA 02111-1307 USA
#

import compiler
import compiler.ast
import logging
import parser

import rete

class RuleCompiler(object):
    def __init__(self, rule_engine, rete):
        self.rule_engine = rule_engine
        self.rete = rete
        
        self.logger = logging.getLogger('PyRete')
        
    def compile(self, function):
        """
        Compile a Python function into a Rete Network.
        
        Raises an Exception if the function is not declared as a Rule or if
        there are any Syntax Errors in the Rule definition. Compilation is
        done in two parts: macro-expansion and Rete Network construction.
        """
        
        self.logger.info("compile(function = [%s])", function.func_name)
        try:
            source_code = function._source
            
        except AttributeError:
            raise Exception("'%s' is not a Rule!" % (function.func_name))
            
        ast = compiler.parse(source_code)
        
        # First phase: Macro expansion
        # Apply all available macro expansions on the Rule.
        macroExpander = MacroExpander()
        compiler.walk(ast, macroExpander)
        ast = macroExpander.ast
        
        # Second phase: Rete Network construction
        reteConstructor = ReteConstructor(self.rule_engine, self.rete)
        compiler.walk(ast, reteConstructor)
        
        
class MacroExpander(object):
    def __init__(self):
        self.ast = None
        
    def visitFunction(self, function):
        """
        Apply all available expansions on function.
        """
        
        ast = function
        for expander in [f for f in dir(self) if f.endswith("_expansion")]:
            ast = getattr(self, expander)(ast)
            
        self.ast = ast
        
    def exists_expansion(self, node):
        """
        Exists(<ce>+) -> Not(Not(<ce>+))
        """
        
        nodesToVisit = [node]
        while len(nodesToVisit) > 0:
            curr = nodesToVisit.pop(0)
            for child in curr.getChildNodes():
                nodesToVisit.append(child)
                
            if isinstance(curr, compiler.ast.CallFunc) and \
               curr.node.name == "Exists":
                curr.node.name = "Not"                                          # Exists(<ce>+) -> Not(<ce>+)
                curr.args = [compiler.ast.CallFunc(compiler.ast.Name("Not"),    # Not(<ce>+) -> Not(Not(<ce>+))
                                                   curr.args)]
                                                   
        return node
        
    def forall_expansion(self, node):
        """
        Forall(<ce>, <ce>+) -> Not(<ce>, Not(<ce>+))
        """
        
        nodesToVisit = [node]
        while len(nodesToVisit) > 0:
            curr = nodesToVisit.pop(0)
            for child in curr.getChildNodes():
                nodesToVisit.append(child)
                
            if isinstance(curr, compiler.ast.CallFunc) and \
               curr.node.name == "Forall":
                if len(curr.args) < 2:
                    raise SyntaxError("Forall-CE must have at least 2 Object-Pattern-CEs!")
                    
                curr.node.name = "Not"                                          # Forall(<ce>, <ce>+) -> Not(<ce>, <ce>+)
                curr.args = [curr.args[0],                                      # Not(<ce>, <ce>+) -> Not(<ce>, Not(<ce>+))
                             compiler.ast.CallFunc(compiler.ast.Name("Not"),
                                                   curr.args[1:])]
                                                   
        return node
        
        
class ReteConstructor(object):
    def __init__(self, rule_engine, rete):
        self.rule_engine = rule_engine
        self.rete = rete
        
        self._notStack = 0
        self.alphaNodes = {}
        self.alphaMemory = {}
        self.joinNodes = {}
        self.joinNode = None
        self.previousJoinNode = None
        self.production_node = None
        
        self.rule = { "variables": {}, }
        self.invert = { "==" : "==",
                        "!=" : "!=",
                        "<"  :  ">",
                        "<=" : ">=",
                        ">"  : "<",
                        ">=" : "<=", }
                        
        self._current_class, self._current_fact = None, None
        
        self.logger = logging.getLogger('PyRete')
        
    def _add_to_test_network(self, node, type):
        """
        Add an AlphaNode to the Rete Network.
        
        If <node> is the first one of <type> it's added directly to the 
        object type node, otherwise it's added to the last AlphaNode
        """
        
        self.logger.debug("ReteConstructor._add_to_test_network(node = %s, type = %s)", node, type)
        if type not in self.alphaMemory:
            self.alphaMemory[type] = rete.Memory()
            
        if self.alphaNodes[type] is None:
            self.rete._add_to_root(node, type)
            
        else:
            prev = self.alphaNodes[type]
            self.rete._connect_nodes(prev, node)
            
        self.alphaNodes[type] = node
        
    def visitFunction(self, function):
        """
        Handles the construction of a Rule.
        """
        
        self.production_node = rete.ProductionNode(rule = function.name)
        self.production_node.doc = function.doc
        
        # Function arguments are used to specify a Rule's Fact bindings.
        # Bindings are neccessary if a Fact is to be referenced in the RHS
        # of the Rule. The default value must be the Fact object's type
        # (class name).
        arguments, defaults = function.argnames, function.defaults
        for argument in reversed(arguments):
            type = None
            if len(defaults) > 0:
                type = defaults.pop(-1).name
                
            if type:
                self.alphaNodes[type] = None
                self.rule[argument] = type
                
        decl_part_parsed, if_part_parsed = False, False
        for node in function.code:
            # Assignments
            if isinstance(node, compiler.ast.Assign):
                decl_part_parsed = True
                if if_part_parsed:
                    raise SyntaxError("Declarations must be done before the LHS of a rule!")
                    
                compiler.walk(node, self)
                
            # If-Then statement
            elif isinstance(node, compiler.ast.If):
                if if_part_parsed:
                    raise SyntaxError("Cannot have more than one if statement in a rule!")
                    
                if_part_parsed = True
                compiler.walk(node, self)
                
            # Everything else is disallowed!
            else:
                raise SyntaxError("%s is not allowed!" % (node))
                
        # Add variables and their bindings to the production node
        for variable in self.rule["variables"]:
            self.production_node.bindings[variable] = self.rule["variables"][variable][0]
            
        # Tie it all together by connecting the last node in this Rule's
        # Rete Network to the Production Node.
        self.rete._connect_nodes(self.previousJoinNode, self.production_node)
        self.rete._add_production(function.name, self.production_node)
        
    def visitAssign(self, assign):
        """
        Handles assignment of Rule properties.
        """
        
        for target in assign.nodes:
            if target.name == "salience" and \
               type(assign.expr.value) in [int, float]:
                self.production_node.salience = assign.expr.value
                
            elif target.name == "autofocus" and \
               type(assign.expr.value) == bool:
                self.production_node.autofocus = assign.expr.value
                
            else:
                raise SyntaxError("Assigning %s to %s is not allowed!" % (repr(assign.expr.value), target.name))
                
                
    def wrapAST(self, ast):
        return compiler.ast.Module(None, ast)
        
    def visitIf(self, ifThenStatement):
        """
        Handles the LHS of a Rule.
        """
        
        if len(ifThenStatement.tests)> 1:
            raise SyntaxError("Else-If branches are not allowed in a Rule!")
            
        elif ifThenStatement.else_ is not None:
            raise SyntaxError("An Else branch is not allowed in a Rule!")
            
        else:
            expr, consequence = ifThenStatement.tests[0]
            
            if isinstance(expr, compiler.ast.And):
                for node in expr.nodes:
                    compiler.walk(node, self)
                    
            elif isinstance(expr, compiler.ast.Compare):
                compiler.walk(expr, self)
                
            elif isinstance(expr, compiler.ast.CallFunc):
                compiler.walk(expr, self)
                
            # Check for unsupported constructs. the RHSValidator raises a
            # SyntaxError exception if any unsupported functionality is
            # found.
            compiler.walk(consequence, RHSValidator())
            self.production_node.rhs = self.wrapAST(consequence)
            
    def visitCallFunc(self, callFunc):
        """
        Handles conditional elements.
        
        The Rete Network implementation can handle three different types of
        conditional elements: Object-patterns, Tests and Nots.
        """
        
        if callFunc.node.name == "Not":
            self._notStack += 1
            
            for arg in callFunc.args:
                compiler.walk(arg, self)
                
            self._notStack -= 1
            
        elif callFunc.node.name == "Test":
            testNode = rete.TestNode(callFunc.args)
            
            if not self.previousJoinNode:
                raise Exception("A Test-CE cannot be at the top of the Rete Network!")
                
            else:
                # the AlphaMemory with the current JoinNode
                self.rete._connect_nodes(mem, testNode, "right")
                
                # and the last JoinNode Memory with the current JoinNode
                self.rete._connect_nodes(self.previousJoinNode, testNode, "left")
                
                self.previousJoinNode = rete.Memory()
                self.rete._connect_nodes(testNode, self.previousJoinNode)
                
                
        else: # This is an Object-Pattern-CE
            self._current_fact = callFunc.node.name
            self._current_class = self.rule[self._current_fact]
            
            self.joinNode = rete.DummyJoinNode(self._current_fact)
            
            # Construct test nodes for all constraints in this CE.
            for arg in callFunc.args:
                compiler.walk(arg, self)
                
            # Join the last of the AlphaNodes with the AlphaMemory,
            if self._current_class not in self.alphaMemory:
                self.alphaMemory[self._current_class] = rete.Memory()
                
            mem = self.alphaMemory[self._current_class]
            if self.alphaNodes[self._current_class] is None:
                self.rete._add_to_root(mem, self._current_class)
            else:
                self.rete._connect_nodes(self.alphaNodes[self._current_class], mem)
                
            if not self.previousJoinNode:
                if self._notStack > 0:
                    raise Exception("A Not-CE cannot be at the top of the Rete Network!")
                    self.previousJoinNode = rete.NotNode()
                    
                else:
                    self.previousJoinNode = rete.LeftInputAdapter(self._current_fact)
                    self.rete._connect_nodes(mem, self.previousJoinNode, "left")
                    
            else:
                if self._notStack > 0:
                    notNode = rete.NotNode()
                    for test in self.joinNode.tests:
                        notNode.add_test(*test)
                    self.joinNode = notNode
                    
                # the AlphaMemory with the current JoinNode
                
                self.rete._connect_nodes(mem, self.joinNode, "right")
                
                # and the last JoinNode Memory with the current JoinNode
                self.rete._connect_nodes(self.previousJoinNode, self.joinNode, "left")
                
                self.previousJoinNode = rete.Memory()
                self.rete._connect_nodes(self.joinNode, self.previousJoinNode)
                
    def visitAnd(self, nodes):
        for node in nodes:
            compiler.walk(node, self)
            
            
    def visitCompare(self, compare):
        """
        Handles constraints in conditional elements.
        
        Constraints are compiled into either AlphaNodes (Tests) or BetaNodes
        (Joins).
        """
        
        # Figure out what type of comparison this is. p1 and p2 are the
        # values compared, p1t and p2t are their types and rel is the
        # comparison operator.
        fact_cls = eval(self._current_class, globals(), self.rule_engine.env)
        p1, p1t, rel, p2, p2t = None, None, None, None, None
        try:
            p1 = compare.expr.value
            p1t = "Constant"
        except:
            try:
                p1 = compare.expr.name
                if p1 in fact_cls.__slots__:
                    p1t = "Attribute"
                else:
                    p1t = "Variable"
            except:
                raise SyntaxError("---")
                
        try:
            p2 = compare.ops[0][1].value
            p2t = "Constant"
        except:
            try:
                p2 = compare.ops[0][1].name
                if p2 in fact_cls.__slots__:
                    p2t = "Attribute"
                else:
                    p2t = "Variable"
            except:
                raise SyntaxError("---")
                
        rel = compare.ops[0][0]
        
        # Testing two constants doesn't make sense, insult the user and
        # refuse to compile this rule!
        if p1t == "Constant" and p2t == "Constant":
            raise SyntaxError("Idiot! The test %s %s %s is pointless! Take it away or I won't compile your stupid rule." % (p1, rel, p2))
            
            
        # Variables and Constants
        elif p1t == "Variable" and p2t == "Constant":
            if p1 in self.rule["variables"].keys():
                a = rete.make_ConstantTestNode(self.rule["variables"][p1][0], rel, p2)
                self._add_to_test_network(a, self._current_class)
                
            else:
                raise SyntaxError("%s is unbound!" % (p1))
                
        elif p1t == "Constant" and p2t == "Variable":
            if p2 in self.rule["variables"].keys():
                # The function make_ConstantTestNode assumes that the
                # comparison is written with the Variable first so we'll
                # have to switch places of p1 and p2 and also invert the
                # comparison operator.
                a = rete.make_ConstantTestNode(self.rule["variables"][p2][0], self.invert[rel], p1)
                self._add_to_test_network(a, self._current_class)
                
            else:
                raise SyntaxError("%s is unbound!" % (p2))
                
        # Attributes and Constants
        elif p1t == "Attribute" and p2t == "Constant":
            a = rete.make_ConstantTestNode(p1, rel, p2)
            self._add_to_test_network(a, self._current_class)
            
        elif p1t == "Constant" and p2t == "Attribute":
            # The function make_ConstantTestNode assumes that the
            # comparison is written with the Variable first so we'll
            # have to switch places of p1 and p2 and also invert the
            # comparison operator.
            a = rete.make_ConstantTestNode(p2, self.invert[rel], p1)
            self._add_to_test_network(a, self._current_class)
            
        # Variables and Attributes
        elif p1t == "Variable" and p2t == "Attribute":
            if p1 in self.rule["variables"].keys():
                if self.rule["variables"][p1][1] == self._current_class:
                    # If the variable in p1 is bound to an attribute in
                    # the current_class we handle the test in an AlphaNode
                    if rel == "==":
                        # If the test is == we use a BindingsTestNode ...
                        a = rete.make_BindingsTestNode(self.rule["variables"][p1][0], p2)
                        self._add_to_test_network(a, self._current_class)
                        
                    else:
                        # ... otherwise we'll use an AttributeTestNode.
                        a = rete.make_AttributeTestNode(self.rule["variables"][p1][0], rel, p2)
                        self._add_to_test_network(a, self._current_class)
                        
                else:
                    # If it's bound to an attribute in another class we'll
                    # have to use a JoinNode
                    jkey = self.rule["variables"][p1][2], self._current_fact
                    if jkey in self.joinNodes:
                        # If a join between these two fact classes already
                        # exist we'll add a test to it ...
                        self.joinNodes[jkey].add_test(self.rule["variables"][p1][2], self.rule["variables"][p1][3], rel, self._current_fact, p2)
                        
                    else:
                        # ... otherwise we'll construct a new JoinNode
                        self.joinNodes[jkey] = rete.JoinNode(self.rule["variables"][p1][2], self.rule["variables"][p1][3], rel, self._current_fact, p2)
                        self.joinNode = self.joinNodes[jkey]
                        
            elif rel == "==":
                self.rule["variables"][p1] = "%s.%s" % (self._current_fact, p2), self._current_class, self._current_fact, p2
                self.logger.debug("ProductionNode.env['%s'] = %s", p1, self.rule["variables"][p1][0])
                
            else:
                raise SyntaxError("%s is unbound!" % (p2))
                
        elif p1t == "Attribute" and p2t == "Variable":
            if p2 in self.rule["variables"].keys():
                if self.rule["variables"][p2][1] == self._current_class:
                    # If the variable in p2 is bound to an attribute in
                    # the current_class we handle the test in an AlphaNode
                    if rel == "==":
                        # If the test is == we use a BindingsTestNode ...
                        a = rete.make_BindingsTestNode(self.rule["variables"][p2][3], p1)
                        self._add_to_test_network(a, self._current_class)
                        
                    else:
                        # ... otherwise we'll use an AttributeTestNode.
                        a = rete.make_AttributeTestNode(p1, rel, self.rule["variables"][p2][0])
                        self._add_to_test_network(a, self._current_class)
                        
                else:
                    # If it's bound to an attribute in another class we'll
                    # have to use a JoinNode
                    jkey = self.rule["variables"][p2][2], self._current_fact
                    if jkey in self.joinNodes:
                        # If a join between these two fact classes already
                        # exist we'll add a test to it ...
                        self.joinNodes[jkey].add_test(self.rule["variables"][p2][2], self.rule["variables"][p2][3], rel, self._current_fact, p1)
                        
                    else:
                        # ... otherwise we'll construct a new JoinNode
                        self.joinNodes[jkey] = rete.JoinNode(self.rule["variables"][p2][2], self.rule["variables"][p2][3], rel, self._current_fact, p1)
                        self.joinNode = self.joinNodes[jkey]
                        
            elif rel == "==":
                self.rule["variables"][p2] = "%s.%s" % (self._current_fact, p1), self._current_class, self._current_fact, p1
                self.logger.debug("ProductionNode.env['%s'] = %s", p2, self.rule["variables"][p2][0])
                
            else:
                raise SyntaxError("%s is unbound!" % (p2))
                
        # Variables and Variables
        elif p1t == "Variable" and p2t == "Variable":
            if p1 in self.rule["variables"].keys() and \
               p2 in self.rule["variables"].keys():
                raise SyntaxError("Cannot compare two variables!")
                ##if self.rule["variables"][p1][1] == self.rule["variables"][p2][1]:
                ##    if rel == "==":
                ##        # If the test is == we use a BindingsTestNode ...
                ##        a = rete.make_BindingsTestNode(self.rule["variables"][p1][3], self.rule["variables"][p2][3])
                ##        self._add_to_test_network(a, self.rule["variables"][p1][1])
                ##        
                ##    else:
                ##        # ... otherwise we'll use an AttributeTestNode.
                ##        a = rete.make_AttributeTestNode(self.rule["variables"][p1][3], rel, self.rule["variables"][p2][3])
                ##        self._add_to_test_network(a, self.rule["variables"][p1][1])
                ##        
                ##else:
                ##    jkey = self.rule["variables"][p1][2], self.rule["variables"][p2][2]
                ##    if jkey in self.joinNodes:
                ##        self.joinNodes[jkey].add_test(self.rule["variables"][p1][2], self.rule["variables"][p1][3], rel, self.rule["variables"][p2][2], self.rule["variables"][p2][3])
                ##        
                ##    else:
                ##        # ... otherwise we'll construct a new JoinNode
                ##        self.joinNodes[jkey] = rete.JoinNode(self.rule["variables"][p1][2], self.rule["variables"][p1][3], rel, self.rule["variables"][p2][2], self.rule["variables"][p2][3])
                ##        self.joinNode = self.joinNodes[jkey]
                ##        
            elif p1 not in self.rule["variables"].keys():
                raise SyntaxError("%s is unbound!" % (p1))
                
            elif p2 not in self.rule["variables"].keys():
                raise SyntaxError("%s is unbound!" % (p2))
                
        # Attributes and Attributes
        elif p1t == "Attributes" and p2t == "Attributes":
            if rel == "==":
                # If the test is == we use a BindingsTestNode ...
                a = rete.make_BindingsTestNode(p1, p2)
                self._add_to_test_network(a, self._current_class)
                
            else:
                # ... otherwise we'll use an AttributeTestNode.
                a = rete.make_AttributeTestNode(p1, rel, p2)
                self._add_to_test_network(a, self._current_class)
                
                
class RHSValidator(object):
    """
    RHSValidator is a compiler.visitor object that raises a SyntaxError
    Exception if an unsupported language construct is detected.
    
    At the moment this is a very simplistic (maybe too restrictive)
    implementation.
    """
    
    def visitFunction(self, func):
        """
        This is required in order to prevent visiting the constructs in a
        function definition.
        """
        pass
        
    def visitClass(self, cls):
        raise SyntaxError("Class definitions are not allowed in the RHS of a Rule!")
        
    def visitGenExpr(self, genExpr):
        raise SyntaxError("Generator expressions are not allowed in the RHS of a Rule!")
        
    def visitGenExprFor(self, genExpr):
        raise SyntaxError("Generator expressions are not allowed in the RHS of a Rule!")
        
    def visitGenExprIf(self, genExpr):
        raise SyntaxError("Generator expressions are not allowed in the RHS of a Rule!")
        
    def visitGenExprInner(self, genExpr):
        raise SyntaxError("Generator expressions are not allowed in the RHS of a Rule!")
        
    def visitGlobal(self, glob):
        raise SyntaxError("Global variables are not allowed in the RHS of a Rule!")
        
    def visitIf(self, ifStmt):
        raise SyntaxError("If statements are not allowed in the RHS of a Rule!")
        
    def visitImport(self, importStmt):
        raise SyntaxError("Imports are not allowed in the RHS of a Rule!")
        
    def visitReturn(self, returnStmt):
        raise SyntaxError("Return statements are not allowed in the RHS of a Rule!")
        
    def visitTryExcept(self, tryExceptStmt):
        raise SyntaxError("Try, Except and Finally are not allowed in the RHS of a Rule!")
        
    def visitTryFinally(self, tryFinallyStmt):
        raise SyntaxError("Try, Except and Finally are not allowed in the RHS of a Rule!")
        
    def visitYield(self, yieldStmt):
        raise SyntaxError("Yield statements are not allowed in the RHS of a Rule!")
        
        